from tube_generate import *
import numpy as np
import matplotlib.pyplot as plt
from plyfile import PlyData, PlyElement
import math
from scipy.stats import gaussian_kde

def calculate_circle(p1, p2, p3):
    # 转换为二维坐标（假设在XY平面）
    p1_2d = p1[:2]
    p2_2d = p2[:2]
    p3_2d = p3[:2]
    
    # 检查共线性
    if np.linalg.matrix_rank([p2_2d - p1_2d, p3_2d - p1_2d]) < 2:
        return None, None
    
    # 构建2x2系数矩阵
    A = np.array([
        [2*(p2_2d[0] - p1_2d[0]), 2*(p2_2d[1] - p1_2d[1])],
        [2*(p3_2d[0] - p2_2d[0]), 2*(p3_2d[1] - p2_2d[1])]
    ])
    
    # 构建常数项
    b = np.array([
        p2_2d[0]**2 - p1_2d[0]**2 + p2_2d[1]**2 - p1_2d[1]**2,
        p3_2d[0]**2 - p2_2d[0]**2 + p3_2d[1]**2 - p2_2d[1]**2
    ])
    
    try:
        center_2d = np.linalg.solve(A, b)
        radius = np.linalg.norm(center_2d - p1_2d)
        # 恢复z坐标为三点平均值
        z_center = (p1[2] + p2[2] + p3[2]) / 3
        return np.array([center_2d[0], center_2d[1], z_center]), radius
    except np.linalg.LinAlgError:
        return None, None

def detect_circle_radom_3_points(points,retry_num=3):
    circle_points=[]
    circle_r=[]

    for j in range(retry_num):
        # 打乱点集顺序
        shuffled_points = np.random.permutation(points)
        
        i = 0
        while i < len(shuffled_points) - 2:
            # 取连续三个不重复的点
            triplet = shuffled_points[i:i+3]
            
            # 计算圆心和半径
            center, radius = calculate_circle(*triplet)
            
            if center is not None and radius is not None:
                circle_points.append(center)

                circle_r.append(radius)
            
            # 跳过已处理的点
            i += 3

    return circle_points, circle_r

def detect_circle(circle_points,circle_r):
    if len(circle_points) > 0:
        # 过滤异常半径（超过1.5倍IQR）
        radii = np.array(circle_r)
        q75, q25 = np.percentile(radii, [75, 25])
        iqr = q75 - q25
        mask = (radii > q25 - 1.5*iqr) & (radii < q75 + 1.5*iqr)
        
        # 计算加权平均值（根据半径密度加权）
        filtered_points = np.array(circle_points)[mask]
        filtered_radii = radii[mask]
        
        # 使用核密度估计找到最密集区域
        kde = gaussian_kde(filtered_radii)
        xvals = np.linspace(min(filtered_radii), max(filtered_radii), 100)
        kde_vals = kde(xvals)
        best_radius = xvals[np.argmax(kde_vals)]
        
        # 找到对应半径的所有圆心
        candidate_centers = filtered_points[np.abs(filtered_radii - best_radius) < iqr*0.5]
        
        # 取三维坐标中位数作为最终圆心
        final_center = np.median(candidate_centers, axis=0)
        final_radius = np.median(filtered_radii)

        radius_std = np.std(filtered_radii - final_radius)
           
    return final_center,final_radius,radius_std

def detect_tube(ply_file,choose_depth=10):
    # 读取PLY文件
    ply_data = PlyData.read(ply_file)
    vertices = ply_data['vertex']
    
    # 提取坐标数据
    x = vertices['x']
    y = vertices['y']
    z = vertices['z']

    # 计算深度范围
    min_depth = np.min(z)
    max_depth = np.max(z)
    depth_range = max_depth - min_depth

    # 筛选深度在指定范围内的点
    x_list = []
    y_list = []
    z_list = []
    
    # 动态计算bins数量（每个区间宽度为choose_depth）
    num_bins = int(np.ceil(depth_range / choose_depth))
    
    # 生成直方图
    hist, bin_edges = np.histogram(z, bins=num_bins, range=(min_depth, max_depth))
    max_bin_index = np.argmax(hist)
    z_min = bin_edges[max_bin_index]
    z_max = bin_edges[max_bin_index + 1]

    x_list = []
    y_list = []
    z_list = []

    # 筛选最密集区间内的点
    for i in range(len(z)):
        if z[i] >= z_min and z[i] <= z_max:
            x_list.append(x[i])
            y_list.append(y[i])
            z_list.append(z[i])

    # 在筛选点之后添加以下代码
    points = np.column_stack((x_list, y_list, z_list))
    circle_points = []
    circle_r = []

    circle_points,circle_r = detect_circle_radom_3_points(points)

    circle = detect_circle(circle_points,circle_r)  

    colors = []
    for depth in z:
        # 将深度映射到颜色索引 (0-255)
        color_idx = int(((depth - min_depth) / depth_range) * (len(color_map)-1))
        colors.append(color_map[color_idx])
    
    # 转换为0-1范围的RGB值
    colors = np.array(colors) / 255.0

    # 创建3D可视化
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 绘制原始点云
    ax.scatter(x, y, z, c=colors, s=1, marker='o', label='原始点云')
    
    # 绘制筛选后的点云（红色突出显示）
    # ax.scatter(x_list, y_list, z_list, c="black", s=1, marker='o', label='筛选点')

    # 绘制筛选后的圆心和半径
    # if len(circle_points) > 0:
    #     for center, radius in zip(circle_points, circle_r):
    #         # 在XY平面生成圆环
    #         theta = np.linspace(0, 2*np.pi, 100)
    #         x_circle = center[0] + radius * np.cos(theta)
    #         y_circle = center[1] + radius * np.sin(theta)
            
    #         # 使用圆心Z坐标的平均值作为高度
    #         z_circle = np.full_like(x_circle, center[2])
            
    #         ax.plot(x_circle, y_circle, z_circle, color='red', linewidth=2, 
    #                 label='检测圆环' if center is circle_points[0] else "")
    
    print(circle)

    if len(circle[0]) and circle[1]:
        # 生成更密集的点用于最终圆环
        theta_final = np.linspace(0, 2*np.pi, 500)
        x_final = circle[0][0] + circle[1] * np.cos(theta_final)
        y_final = circle[0][1] + circle[1] * np.sin(theta_final)
        z_final = np.full_like(theta_final, circle[0][2])
        
        # 使用蓝色粗线突出显示
        if circle[2]<20:
            ax.plot(x_final, y_final, z_final, color='blue', linewidth=4, 
                linestyle='--', label='最终检测圆环')
        else:
            ax.plot(x_final, y_final, z_final, color='black', linewidth=4, 
                linestyle='--', label='最终检测圆环')
        
        # 添加半径标注
        ax.text(circle[0][0], circle[0][1], circle[0][2], 
               f'R={circle[1]:.2f}mm', color='blue', fontsize=12)
        
    # 设置坐标轴标签
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    
    plt.show(block=True)  # 添加block=True参数保持窗口显示

if __name__ == "__main__":
    detect_tube('tube.ply')